import struct
import serial
import sys
import time

FIRMWARE_SIZE_PER_PACK = 512
def invert_uint8(data):
	temp = 0
	for i in range(8):
		if data & (1<<i) :
			temp = temp | (1<<(7-i))
	return temp


def inver_uint16(data):
	temp = 0
	for i in range(16):
		if data & (1 << i):
			temp = temp | (1 << (15 - i))
	return temp


def crc16_modbus(buffer):
	crc_in = 0xFFFF
	poly = 0x8005
	temp = 0
	for item in buffer:
		temp = item
		temp = invert_uint8(temp)
		crc_in = crc_in ^ (temp << 8)
		for i in range(8):
			if crc_in & 0x8000:
				crc_in = (crc_in << 1) ^ poly
			else:
				crc_in = crc_in << 1
	crc_in = inver_uint16(crc_in)
	return crc_in


def print_buf(buffer):
	for i in buffer:
		print('%02X'%i + ' ', end='')
	print('')


def make_frame(ctrl, buffer) :
	prefix = b'\x55\xaa'
	ctrl = struct.pack('<H', ctrl)
	length = struct.pack('<H', len(buffer))
	if len(buffer) > 0 :
		body = ctrl+length+buffer
	else:
		body = ctrl+length
	crc16 = struct.pack('<H', crc16_modbus(body))
	return prefix+body+crc16


def check_frame(buffer) :
	prefix = b'\x55\xaa'
	length = len(buffer)
	if length < 8:
		return []
	for offset in range(length - 8 + 1):
		if prefix != buffer[offset:offset+2]:
			continue
		body_len = struct.unpack('<H', buffer[offset+4:offset+6])[0]
		crc16 = struct.pack('<H', crc16_modbus(buffer[offset+2:offset+2+4+body_len]))
		if crc16 != buffer[offset+2+4+body_len:offset+2+4+body_len+2]:
			continue
		else:
			return buffer[offset:offset+2+4+body_len+2]
	return []


def progress_bar(value, total):
	bar_length = 50
	progress = int(value / total * bar_length)
	print('\r', end='')
	print('[', end='')
	if progress == 0:
		print('>', end='')
		for i in range(bar_length-1):
			print(' ', end='')
	elif progress != bar_length:
		for i in range(progress):
			print('=', end='')
		print('>', end='')
		for i in range(bar_length-1-progress):
			print(' ', end='')
	else:
		for i in range(bar_length):
			print('=', end='')
	print(']', end='')
	print(' %3.1f%%'%(value / total * 100), end='')


def ser_send(ctrl, buf):
	send_buf = make_frame(ctrl, buf)
	ser.write(send_buf)
	return send_buf


def ser_read(timeout):
	timeout = timeout/1000 + time.time()
	temp_buf = bytes(0)
	while time.time() < timeout:
		if ser.in_waiting:
			temp_buf = temp_buf + ser.read(1)
		recv_buf = check_frame(temp_buf)
		if len(recv_buf) != 0:
			return recv_buf
	return []


def file_data_read(offset, size):
	file.seek(offset, 0)
	return file.read(size)


def connect():
	ctrl = 0x0000
	file.seek(0,2)
	file_size = file.tell()
	file_size = struct.pack('<i', file_size)
	ser_send(ctrl, file_size)
	recv_buf = ser_read(3000)
	if len(recv_buf) != 0:
		recv_buf = recv_buf[2:-2]
		reply_ctrl = struct.unpack('<H', recv_buf[0:2])[0]
		if reply_ctrl != (ctrl|0x8000):
			return 'timeout'
		data_len = struct.unpack('<H', recv_buf[2:4])[0]
		if data_len == 1:
			data = recv_buf[4:]
			if data == b'\xac' :
				return 'accept'
		else:
				return 'refuse'
	else:
		return 'timeout'


def send_firmware():
	ctrl = 0x0001
	file.seek(0,2)
	file_size = file.tell()
	remain = file_size
	write_pos = 0
	count = 0
	file.seek(0, 0)
	print('Transfer progress:')
	progress_bar(file_size-remain, file_size)
	while True:
		if remain > FIRMWARE_SIZE_PER_PACK:
			read_size = FIRMWARE_SIZE_PER_PACK
		else:
			read_size = remain
		file.seek(write_pos, 0)
		firmware_data = file.read(read_size)
		data = struct.pack('<LH', write_pos, read_size) + firmware_data
		ser_send(ctrl, data)
		recv_buf = ser_read(3000)
		if len(recv_buf) != 0:
			recv_buf = recv_buf[2:-2]
			reply_ctrl = struct.unpack('<H', recv_buf[0:2])[0]
			if reply_ctrl != (ctrl|0x8000):
				return 'refuse'
			data_len = struct.unpack('<H', recv_buf[2:4])[0]
			if data_len == 1:
				data = recv_buf[4:]
				if data == b'\xac':
					remain -= read_size
					progress_bar(file_size - remain, file_size)
					if remain == 0:
						print('')
						return 'accept'
					else:
						count = 0
						write_pos += read_size
						continue
				else:
					return 'refuse'
		else:
			count += 1
			if count < 4:
				continue
			else:
				return "refuse"

def verify_firmware():
	ctrl = 0x0002
	ser_send(ctrl, [])
	recv_buf = ser_read(3000)
	if len(recv_buf) != 0:
		recv_buf = recv_buf[2:-2]
		reply_ctrl = struct.unpack('<H', recv_buf[0:2])[0]
		if reply_ctrl != (ctrl|0x8000):
			return 'timeout'
		data_len = struct.unpack('<H', recv_buf[2:4])[0]
		if data_len == 1:
			data = recv_buf[4:]
			if data == b'\xac':
				return 'accept'
			else:
				return 'refuse'
	else:
		return 'timeout'


def update_completed():
	ctrl = 0x0003
	ser_send(ctrl, [])
	recv_buf = ser_read(3000)
	if len(recv_buf) != 0:
		recv_buf = recv_buf[2:-2]
		reply_ctrl = struct.unpack('<H', recv_buf[0:2])[0]
		if reply_ctrl != (ctrl|0x8000):
			return 'timeout'
		data_len = struct.unpack('<H', recv_buf[2:4])[0]
		if data_len == 1:
			data = recv_buf[4:]
			if data == b'\xac':
				return 'accept'
			else:
				return 'refuse'
	else:
		return 'timeout'


def main():
	retry = 0
	connect_retry = 0
	state = 'CONNECT'
	print('start update ...')
	while True:
		print(state)
		if state == 'CONNECT':
			ret = connect()
			if ret == 'accept':
				state = 'SEND_FIRMWARE'
			else:
				connect_retry += 1
				if connect_retry >= 3:
					return 'Device cannot be connected normally.'
		elif state == 'SEND_FIRMWARE':
			ret = send_firmware()
			if ret == 'accept':
				state = 'VERIFY_FIRMWARE'
			elif ret == 'refuse':
				state = 'CONNECT'
		elif state == 'VERIFY_FIRMWARE':
			ret = verify_firmware()
			if ret == 'accept':
				retry = 0
				state = 'UPDATE_COMPLETED'
			elif ret == 'refuse':
				retry = 0
				state = 'CONNECT'
			elif ret == 'timeout':
				retry += 1
				if retry > 3:
					retry = 0
					state = 'CONNECT'
		elif state == 'UPDATE_COMPLETED':
			ret = update_completed()
			if ret == 'accept':
				break
			elif ret == 'refuse':
				retry = 0
				state = 'CONNECT'
			elif ret == 'timeout':
				retry += 1
				if retry > 3:
					retry = 0
					state = 'CONNECT'
	if connect_retry < 3:
		print('Update completed.')
	else:
		print('Update fail.')



if __name__ == '__main__':
	if len(sys.argv) < 3:
		print('Two argument expected. (example: python update_tools.py COMx file_path)')
	try:
		print('Open serial', sys.argv[1])
		ser = serial.Serial(sys.argv[1], baudrate=115200)
	except:
		print('Serial port open error.')
		exit(-1)
	try:
		print('Open file', sys.argv[2])
		file = open(sys.argv[2], "rb")
	except:
		print(sys.argv[2], ': No such file or directory.')
		ser.close()
		exit(-1)
	main()
	print('Close file', sys.argv[2])
	file.close()
	print('Close serial', sys.argv[1])
	ser.close()
	exit(0)

